use std::collections::HashMap;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use crate::actor::ActorRef;
use crate::message::DynamicMessage;
use crate::errors::SendError;

#[derive(Clone)]
pub struct BatchConfig {
    pub max_batch_size: usize,
    pub max_batch_interval: Duration,
}

pub struct MessageBatch {
    messages: Vec<DynamicMessage>,
    last_batch_time: Instant,
    config: BatchConfig,
}

impl MessageBatch {
    pub fn new(config: BatchConfig) -> Self {
        Self {
            messages: Vec::new(),
            last_batch_time: Instant::now(),
            config,
        }
    }

    pub fn add_message(&mut self, msg: DynamicMessage) {
        self.messages.push(msg);
    }

    pub fn should_flush(&self) -> bool {
        self.messages.len() >= self.config.max_batch_size
            || self.last_batch_time.elapsed() >= self.config.max_batch_interval
    }

    pub fn take_messages(&mut self) -> Vec<DynamicMessage> {
        let messages = std::mem::take(&mut self.messages);
        self.last_batch_time = Instant::now();
        messages
    }
}

pub struct BatchProcessor {
    batches: HashMap<String, MessageBatch>,
    sender: mpsc::UnboundedSender<(String, Vec<DynamicMessage>)>,
    config: BatchConfig,
}

impl BatchProcessor {
    pub fn new(sender: mpsc::UnboundedSender<(String, Vec<DynamicMessage>)>, config: BatchConfig) -> Self {
        Self {
            batches: HashMap::new(),
            sender,
            config,
        }
    }

    pub async fn add_message(&mut self, target: &ActorRef, msg: DynamicMessage) -> Result<(), SendError> {
        let target_id = target.id().to_string();
        let batch = self.batches.entry(target_id.clone())
            .or_insert_with(|| MessageBatch::new(self.config.clone()));

        batch.add_message(msg);

        if batch.should_flush() {
            self.flush_batch(&target_id).await?;
        }

        Ok(())
    }

    async fn flush_batch(&mut self, target_id: &str) -> Result<(), SendError> {
        if let Some(batch) = self.batches.get_mut(target_id) {
            let messages = batch.take_messages();
            if !messages.is_empty() {
                self.sender.send((target_id.to_string(), messages))
                    .map_err(|_| SendError::MailboxFull)?;
            }
        }
        Ok(())
    }

    pub async fn flush_all(&mut self) -> Result<(), SendError> {
        let target_ids: Vec<String> = self.batches.keys().cloned().collect();
        for target_id in target_ids {
            self.flush_batch(&target_id).await?;
        }
        Ok(())
    }
} 